%matplotlib widget
from utils.traffic_model import SignalController

# Specifies the size of the grid city. Cities are represented by a grid map of city_size*city_size
city_size = 3

# Instantiate the class to generate signal control parameters for a given city size
signal_controller = SignalController(city_size)

# Display the table of correspondence between integer indices and parameter values
signal_controller.show_index_table()

import numpy as np

# Generate [3 * city_size * city_size] integer indices ranging from 0 to 19 using integer random numbers
index_list = np.random.randint(0, 20, 3 * city_size * city_size)

# Convert the integer indices to corresponding parameter values (decode)
signal_info = signal_controller.decode(index_list=index_list, disp_flag=True)

from utils.traffic_model import TrafficSimulator2Malls

# Instantiate simulator
traffic_sim = TrafficSimulator2Malls(
    signal_info=signal_info, num_cars=20, ratio_mall=0.5, seed=0
)

# Run the simulation (display result in animation)
traffic_sim.run(steps=500, animation=True)

# Run simulation (output gif animation)
# *About 1 minute of processing time with the following argument settings.
# *Please try gif output in local environment.
# traffic_sim.run(steps=200, gif="./anim.gif")

import os
import torch


def seed_everything(seed=0):
    os.environ["PYTHONHASHSEED"] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.backends.cudnn.deterministic = True

from amplify import FixstarsClient
from datetime import timedelta

client = FixstarsClient()
client.parameters.timeout = timedelta(milliseconds=5000)  # 5000 msec
# client.token = "xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx"  # If you use Amplify in a local environment or Google Colaboratory, enter your Amplify API token.

import torch.nn as nn


class TorchFM(nn.Module):
    def __init__(self, d: int, k: int) -> None:
        super().__init__()
        self.V = nn.Parameter(torch.randn(d, k), requires_grad=True)
        self.lin = nn.Linear(d, 1)

    def forward(self, x):
        out_1 = torch.matmul(x, self.V).pow(2).sum(1, keepdim=True)
        out_2 = torch.matmul(x.pow(2), self.V.pow(2)).sum(1, keepdim=True)
        out_inter = 0.5 * (out_1 - out_2)
        out_lin = self.lin(x)
        out = out_inter + out_lin
        return out

from torch.utils.data import TensorDataset, DataLoader
from sklearn.model_selection import train_test_split
from typing import Type

import copy


def train(
    X: np.ndarray,
    y: np.ndarray,
    model_class: Type[nn.Module],
    model_params: dict[str, int | float],
    batch_size=1024,
    epochs=3000,
    criterion=nn.MSELoss(),
    optimizer_class=torch.optim.AdamW,
    opt_params={"lr": 1},
    lr_sche_class=None,
    lr_sche_params=None,
):
    X_tensor, y_tensor = (
        torch.from_numpy(X).float(),
        torch.from_numpy(y).float(),
    )
    indices = np.array(range(X.shape[0]))
    indices_train, indices_valid = train_test_split(
        indices, test_size=0.2, random_state=42
    )

    train_set = TensorDataset(X_tensor[indices_train], y_tensor[indices_train])
    valid_set = TensorDataset(X_tensor[indices_valid], y_tensor[indices_valid])
    loaders = {
        "train": DataLoader(train_set, batch_size=batch_size, shuffle=True),
        "valid": DataLoader(valid_set, batch_size=batch_size, shuffle=False),
    }

    model = model_class(**model_params)
    best_model_wts = copy.deepcopy(model.state_dict())
    optimizer = optimizer_class(model.parameters(), **opt_params)
    scheduler = None
    if lr_sche_class is not None:
        scheduler = lr_sche_class(optimizer, **lr_sche_params)
    best_score = 1e18
    for _ in range(epochs):
        losses = {"train": 0.0, "valid": 0.0}

        for phase in ["train", "valid"]:
            if phase == "train":
                model.train()
            else:
                model.eval()

            for batch_x, batch_y in loaders[phase]:
                optimizer.zero_grad()
                out = model(batch_x).T[0]
                loss = criterion(out, batch_y)
                losses[phase] += loss.item() * batch_x.size(0)

                with torch.set_grad_enabled(phase == "train"):
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

            losses[phase] /= len(loaders[phase].dataset)

        with torch.no_grad():
            model.eval()
            if best_score > losses["valid"]:
                best_model_wts = copy.deepcopy(model.state_dict())
                best_score = losses["valid"]
        if scheduler is not None:
            scheduler.step()

    with torch.no_grad():
        model.load_state_dict(best_model_wts)
        model.eval()

    out = model(X_tensor).T[0]
    coef = torch.corrcoef(torch.stack((out, y_tensor)))[0, 1].detach()
    return model, coef.item()

def gen_training_data_flow_onehot(D: int, N0: int, true_func, input_generator):
    X = np.zeros((N0, D))
    y = np.zeros(N0)
    print("")
    print("###################################################")
    print(" Making initial training data")
    print("###################################################")
    for i in range(N0):
        print(f"Generating training data set #{i}.")
        # Randomly determine input vector consisting of (0 or 1)
        x = input_generator()
        # If the same input vector has already exist, regenerate the input values
        is_identical = True
        while is_identical:
            is_identical = False
            for j in range(i + 1):
                if np.all(x == X[j, :]):
                    x = input_generator()
                    is_identical = True
                    break
        # Evaluate input value x with objective function f(x) and copy input-output pair (x,f(x)) to training data
        X[i, :] = x
        y[i] = true_func(x)
        print(f"------------------------")
    return X, y

from amplify import VariableGenerator, solve, one_hot, sum
import matplotlib.pyplot as plt
import sys
import time

# FMQA for non-binary inputs


class FMQA_NB:
    def __init__(self, D, N, N0, k, true_func, client, encoder_decoder) -> None:
        assert N0 < N
        self.D = D
        self.N = N
        self.N0 = N0
        self.k = k
        self.true_func = true_func
        self.client = client
        self.nbins_array = encoder_decoder.nbins_array
        self.encoder = encoder_decoder.encoder
        self.decoder = encoder_decoder.decoder
        self.y = None

    # A member function that repeatedly performs (N-N0)x FMQA based on the training data with adding new training data
    def cycle(self, X, y, log: bool = False) -> np.ndarray:
        # Weight for one-hot constraint
        constraint_weight = max([1, int(np.abs(y).max() + 0.5)])
        print("")
        print("###################################################")
        print(f" Starting FMQA cycles... {constraint_weight=}")
        print("###################################################")

        pred_x = X[0]
        pred_y = 1e18
        self.y = y
        i_sta = 0
        for i in range(self.N - self.N0):
            print(f"FMQA Cycle #{i} ")
            start_time = time.perf_counter()
            try:
                x_hat = self.step(X, y, constraint_weight)
            except RuntimeError:
                sys.exit(f"Unknown error, i = {i}")
            # If an input value identical to the found x_hat already exists in the current training data set, a neighboring value is used as the new x_hat.
            is_identical = True
            while is_identical:
                is_identical = False
                for j in range(i + self.N0):
                    if np.all(x_hat == X[j, :]):
                        # Decode binary vector to integer vector and copy to inputs
                        inputs = self.decoder(x_hat)
                        # Get the surrounding values (inputs is an integer vector)
                        inputs += np.random.randint(-1, 2, len(inputs))
                        for i_inp in range(len(inputs)):
                            if inputs[i_inp] < 0:
                                inputs[i_inp] = 0
                            elif inputs[i_inp] > self.nbins_array[i_inp] - 1:
                                inputs[i_inp] = self.nbins_array[i_inp] - 1
                        # Encode from integer vector to binary vector and copy to x_hat
                        x_hat = self.encoder(inputs)
                        if log:
                            print(f"{i=}, Identical x is found, {x_hat=}")
                        is_identical = True
                        break
            # Evaluate the objective function f() with x_hat
            y_hat = self.true_func(x_hat)
            # Add the input-output pair [x_hat, y_hat] to the training data set
            X = np.vstack((X, x_hat))
            y = np.append(y, y_hat)
            self.y = np.append(self.y, y_hat)
            # Copy the input-output pair to [pred_x, pred_y] when the evaluated value of the objective function updates the minimum value
            print(
                f"FMQA Cycle #{i} finished in {time.perf_counter() - start_time:.2f}s",
                end="",
            )
            if pred_y > y_hat:
                pred_y = y_hat
                pred_x = x_hat
                print(f", variable updated, {pred_y=:.2f}")
            else:
                print()
            print(f"------------------------")
        return pred_x

    # Member function to perform one FMQA cycle
    def step(self, X, y, constraint_weight) -> np.ndarray:
        # Train FM
        start_time = time.perf_counter()
        model, corrcoef = train(
            X,
            y,
            model_class=TorchFM,
            model_params={"d": self.D, "k": self.k},
            batch_size=8,
            epochs=1000,
            criterion=nn.MSELoss(),
            optimizer_class=torch.optim.AdamW,
            # Set scheduler to reduce learning rate with number of epochs
            opt_params={"lr": 0.5},
            lr_sche_class=torch.optim.lr_scheduler.StepLR,
            lr_sche_params={"step_size": 50, "gamma": 0.9},
        )
        print(
            f"FM Training: {time.perf_counter() - start_time:.2f}s (corr:{corrcoef:.2f})"
        )
        start_time = time.perf_counter()
        # Extract FM parameters from the trained FM model
        v, w, w0 = list(model.parameters())
        v = v.detach().numpy()
        w = w.detach().numpy()[0]
        w0 = w0.detach().numpy()[0]
        gen = VariableGenerator()  # Instantiate Variable Generator
        q = gen.array("Binary", self.D)  # Generate a binary variable array
        cost = self.__FM_as_QUBO(
            q, w0, w, v
        )  # Define FM as a QUBO equation from FM parameters
        # Construct one-hot constraints for non-binary variables
        constraint_list = []
        ista = 0
        iend = 0
        for i in range(len(self.nbins_array)):
            iend += self.nbins_array[i]
            constraint_list.append(one_hot(q[ista:iend]))
            ista = iend
        constraints = sum(constraint_list)
        # Add up the objective function and constraints, and pass it to the Fixstars Amplify solver
        model = cost + constraint_weight * constraints
        result = solve(model, self.client)
        print(
            f"QUBO: {time.perf_counter() - start_time:.2f}s (AE execution: {result.execution_time.seconds:.2f}s)"
        )
        if len(result.solutions) == 0:
            raise RuntimeError("No solution was found.")
        q_values = q.evaluate(result.best.values)
        return q_values

    # A function that defines FM as a QUBO equation from FM parameters. As with the previously defined TorchFM class, the formula is written as per the acquisition function form of g(x).
    def __FM_as_QUBO(self, x, w0, w, v):
        lin = w0 + (x.T @ w)
        out_1 = np.array([(x * v[:, i]).sum() ** 2 for i in range(self.k)]).sum()
        # Note that x[j] = x[j]^2 since x[j] is a binary variable in the following equation.
        out_2 = np.array([(x * v[:, i] * v[:, i]).sum() for i in range(self.k)]).sum()
        return lin + (out_1 - out_2) / 2

    # A function to plot the history of i-th objective function evaluations performed within the initial training data construction (blue) and during FMQA cycles (red).
    def plot_history(self):
        assert self.y is not None
        fig = plt.figure(figsize=(6, 4))
        plt.plot(
            [i for i in range(self.N0)],
            self.y[: self.N0],
            marker="o",
            linestyle="-",
            color="b",
        )  # Objective function evaluation values at the time of initial training data generation (random process)
        plt.plot(
            [i for i in range(self.N0, self.N)],
            self.y[self.N0 :],
            marker="o",
            linestyle="-",
            color="r",
        )  # Objective function evaluation values during the FMQA cycles (FMQA cycle process)
        plt.plot(
            [i for i in range(self.N)],
            [self.y[:i].min() for i in range(1, self.N + 1)],
            linestyle="-",
            color="k",
        )  # Update history of minimum objective function values
        plt.xlabel("i-th evaluation of f(x)", fontsize=18)
        plt.ylabel("f(x)", fontsize=18)
        plt.tick_params(labelsize=18)
        return fig

class EncDecOneHot:
    def __init__(self, D: int, nbins_array):
        self.D = D
        self.nbins_array = nbins_array

    def gen_random_input(self):
        ista = 0
        iend = 0
        x = np.zeros(self.D, int)
        for i in range(len(self.nbins_array)):
            iend += self.nbins_array[i]
            idx = np.random.randint(ista, iend, 1)[0]
            x[idx] = 1
            ista = iend
        return x

    def encoder(self, inputs):
        if len(inputs) != len(self.nbins_array):
            raise RuntimeError("inputs should be the same length as nbins_array!")
        x = np.zeros(self.D, int)
        i_offset = 0
        for i in range(len(inputs)):
            x[inputs[i] + i_offset] = 1
            i_offset += self.nbins_array[i]
        return x

    def decoder(self, x):
        bit_locs = np.where(x == 1)[0]
        if len(bit_locs) != len(self.nbins_array):
            raise RuntimeError("bit_locs should be the same length as nbins_array!")
        i_offset = 0
        for i in range(1, len(bit_locs)):
            i_offset += self.nbins_array[i - 1]
            bit_locs[i] -= i_offset
        return bit_locs

import time

# Initialize random seed value
seed_everything()

# Specify the size of the grid city, a [city_size*city_size] lattice map
city_size = 4

# Instantiate the class to generate signal control parameters for a given city size
signal_controller = SignalController(city_size)

# 1D array of bin sizes in one-hot encoding for each variable
nbins_array = np.array(
    [
        [
            signal_controller.nbins_red_duration,
            signal_controller.nbins_green_duration,
            signal_controller.nbins_phase,
        ]
        for i in range(city_size**2)
    ]
).reshape(-1)

# Overall size of binary variables (size of binary decision variables handled by QUBO)
D = nbins_array.sum()
print(f"Number of binary decision variables: {D}")

# Instantiate One-hot encoder-decoder class (EncDecOneHot)
enc_dec_one_hot = EncDecOneHot(D, nbins_array)

# Counter for the number of times the objective function is evaluated
n_eval = 0

num_cars = 100
ratio_mall = 0.5


# Objective function (obtain traffic light control parameters --> run traffic simulation --> calculate objective function value)
def obj_func(x):
    global n_eval
    # Convert from the solution (binary) to integer indices representing the traffic light control parameters (red period, green period, phase)
    index_list = enc_dec_one_hot.decoder(x)

    # In addition, the signal_controller.decode method is used to convert from integer indexes to signal control parameters (real numbers)
    signal_info = signal_controller.decode(index_list)

    # Run simulations with different initialization for n_cases times with the same signal control
    n_cases = 5
    costs = []
    print("Cost convergence: ", end="")
    start_time = time.perf_counter()
    for i in range(n_cases):
        # Instantiate TrafficSimulator2Malls (seed value to determine initial conditions is a random number)
        traffic_sim = TrafficSimulator2Malls(
            signal_info=signal_info,
            num_cars=num_cars,
            ratio_mall=ratio_mall,
            seed=np.random.randint(0, 100),
        )
        # Run the simulation
        traffic_sim.run(steps=20000)
        # Get the negative value of the average speed of all cars in each simulation and append to costs
        costs.append(-traffic_sim.get_latest_mean())
        print(f"{sum(costs) / (i + 1):.2f} ", end="")

    # The average output value of the "n_cases" simulations is the final objective function value.
    cost = sum(costs) / n_cases
    print(
        f"\nEvaluation #{n_eval} finished in {time.perf_counter() - start_time:.2f}s. cost:{cost:.2f}"
    )
    n_eval += 1

    return cost


# N and N0 should be generally increased for meaningful optimization. See the description in Sec. 3.1 above.
N = 3  # Number of times the objective function can be evaluated
N0 = 2  # Number of samples in initial training data
k = 20  # Dimensions of vectors in FM (hyperparameters)

# Start time for optimization
start_time = time.perf_counter()

# Generate initial training data
X, y = gen_training_data_flow_onehot(D, N0, obj_func, enc_dec_one_hot.gen_random_input)

# Instantiate the FMQA class
fmqa_solver = FMQA_NB(
    D=D,
    N=N,
    N0=N0,
    k=k,
    true_func=obj_func,
    client=client,
    encoder_decoder=enc_dec_one_hot,
)

# Execute FMQA cycles
pred_x = fmqa_solver.cycle(X, y)

# Output optimization results
print("###################################################")
print(" Optimal input values and corresponding output")
print("###################################################")
print(f"pred x: {pred_x}")
print(f"pred value: {obj_func(pred_x):.2f}")
print(f"Elapsed time: {time.perf_counter() - start_time:.2f}s")

fig = fmqa_solver.plot_history()

# Same condition as the simulation during optimization
city_size = 4
num_cars = 100
ratio_mall = 0.5
signal_controller = SignalController(city_size)

# Construct integer indices from optimization result pred_x
index_list = enc_dec_one_hot.decoder(pred_x)
print(f"{index_list=}")

# Convert integer indices to signal control parameters (real numbers)
signal_info = signal_controller.decode(index_list)

# Instantiate simulator
traffic_sim = TrafficSimulator2Malls(
    signal_info=signal_info, num_cars=num_cars, ratio_mall=ratio_mall
)

# Run the traffic simulation with the optimized signal control
traffic_sim.run(2000, animation=True, interval=100)

Number of binary decision variables: 960

  ###################################################
  Making initial training data
  ###################################################
  Generating training data set #0.
  Cost convergence: -2.40 -3.41 -3.32 -3.53 -3.95 
  Evaluation #0 finished in 52.88s. cost:-3.95
  ------------------------
  Generating training data set #1.
  Cost convergence: -6.56 -6.06 -6.13 -6.00 -5.50 
  Evaluation #1 finished in 50.06s. cost:-5.50
  ------------------------
  Generating training data set #2.
  Cost convergence: -6.19 -6.58 -6.59 -6.56 -6.49 
  Evaluation #2 finished in 47.57s. cost:-6.49
  ------------------------
  Generating training data set #3.
  Cost convergence: -4.18 -4.01 -4.07 -4.10 -4.13 
  Evaluation #3 finished in 53.65s. cost:-4.13
  ------------------------
  Generating training data set #4.
  Cost convergence: -3.79 -4.32 -4.19 -3.73 -3.90 
  Evaluation #4 finished in 55.38s. cost:-3.90
  ------------------------
  Generating training data set #5.
  Cost convergence: -1.68 -3.16 -3.20 -3.22 -3.06 
  Evaluation #5 finished in 56.32s. cost:-3.06
  ------------------------
  Generating training data set #6.
  Cost convergence: -6.94 -5.61 -5.75 -5.71 -5.90 
  Evaluation #6 finished in 50.48s. cost:-5.90
  ------------------------
  Generating training data set #7.
  Cost convergence: -5.38 -6.08 -5.45 -5.38 -5.43 
  Evaluation #7 finished in 51.79s. cost:-5.43
  ------------------------
  Generating training data set #8.
  Cost convergence: -3.43 -3.91 -4.12 -4.55 -4.85 
  Evaluation #8 finished in 50.39s. cost:-4.85
  ------------------------
  Generating training data set #9.
  Cost convergence: -3.77 -4.18 -4.31 -4.37 -4.08 
  Evaluation #9 finished in 54.33s. cost:-4.08
  ------------------------

Black-Box Optimiztion of Traffic Light Control¶

1. MAS for urban congestion¶

1.1. What is MAS?¶

1.2. Traffic congestion model¶

1.3. MAS traffic simulator¶

2. FMQA program implementation (integer input values)¶

2.1. Random seed initialization¶

2.2. Configuration of Amplify client¶

2.3. Implementing FM with PyTorch¶

2.4. Construction of initial training data¶

2.5. Execution class for FMQA cycle¶

2.6. One-hot encoder/decoder implementation¶

3. Optimization of traffic light control¶

3.1. FMQA execution example¶

3.2. Transition of objective function values during the optimization process¶

3.3. MAS traffic simulation with optimized traffic light settings¶

3.4. Example output from this FMQA sample program¶

4. Summary and practical guidelines¶